#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jun 29 10:10:20 2018

@author: jia.liu
"""

import pandas as pd
import numpy as np

# 运行速度太慢 待优化
def my_dummies(data, col='', split_str=''):
    """
    说明:
        1.可进行常规dummies;
        2.对于某列多字段组合，如'class1|class2|class3', 指定分隔符，进行dummies
    参数:
        data: 类型可为dataframe, 也可为series。若为series，返回的是dummies后的dataframe
        col: 类型为str，若data类型为dataframe，需指定col字段。dataframe中待dummies的列
        split_str: 字符串中使用的分隔符
    """
    if type(data) == pd.core.series.Series:
        srs=data
    else:
        srs = data[col]
    all_dum = []
    for x in srs:
        if split_str != '':
            if type(x)== int:
                x = str(x)
            all_dum.extend(x.split(split_str))
        else:
            all_dum = list(set(srs))
            all_dum.remove(np.nan)
    all_dum = pd.unique(all_dum)
    zero_matrix = np.zeros((len(srs), len(all_dum)))
    dummies = pd.DataFrame(zero_matrix, columns=all_dum)
    for i, col_str in enumerate(srs):
        indices = dummies.columns.get_indexer(str(col_str).split('|'))
        if indices != -1:
            dummies.iloc[i, indices] = 1
    if type(data) == pd.core.series.Series:
        return dummies
    else:
        data = data.drop(col, 1)
        data = data.join(dummies.add_prefix(col+'_'))
        return data

# %%
if __name__ == '__main__':
#    movies = pd.read_csv('/Users/jia.liu/Desktop/yabianliang')
#    aa = my_dummies(movies, 'genres',split_str='|')
#    bb = my_dummies(movies, 'title')
    pass
#    pd.get_dummies(need_series)